import optuna
from optuna.trial import TrialState
from main import setup_parser, merge_config, load_json, NamespaceDict
from trainer import _set_device, _set_random, print_args
import logging
import sys
from learners import get_learner
from utils.data_manager import DataManager
from utils.toolkit import count_parameters
import os
from datetime import datetime
import shutil


def _optuna_train(args, trial):
    args["seed"] = args["seed"][0]
    if args.get("train_seed", None) is None:
        print("========== No train_seed in config, use seed instead. ==========")
        args["train_seed"] = args["seed"]

    init_cls = 0 if args["init_cls"] == args["inc_cls"] else args["init_cls"]
    assert init_cls == 0, "init_cls must be 0 for incremental learning"

    current_time = datetime.now().strftime("%Y%m%d%H%M%S")
    current_time = args["prefix"] if args["prefix"] else current_time
    logfile_path = "logs/{}/{}/{}".format(
        args["learner_name"],
        args["dataset"],
        current_time,
    )
    if not os.path.exists(logfile_path):
        os.makedirs(logfile_path)
    logfilename = os.path.join(
        logfile_path,
        "{}_{}_trial[{}]".format(
            args["prefix"],
            args["seed"],
            trial._trial_id,
        ),
    )
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(filename)s] => %(message)s",
        handlers=[
            logging.FileHandler(filename=logfilename + ".log"),
            logging.StreamHandler(sys.stdout),
        ],
    )

    suffix = args["backbone_type"].split("_")[-1]
    shutil.copy2(args["config"], logfile_path)
    shutil.copy2(f"learners/{args['learner_name']}.py", logfile_path)
    shutil.copy2(f"backbone/vit_{suffix}.py", logfile_path)

    _set_random(args["train_seed"])
    _set_device(args)
    print_args(args)

    ckp_path = "ckps/{}/{}/{}/{}".format(
        args["learner_name"],
        args["backbone_type"],
        args["dataset"],
        current_time,
    )
    args["ckp_path"] = ckp_path if getattr(args, "save_ckp", False) else None

    data_manager = DataManager(
        args["dataset"],
        args["shuffle"],
        args["seed"],
        args["init_cls"],
        args["inc_cls"],
        args,
    )

    args["nb_classes"] = data_manager.nb_classes  # update args
    args["nb_tasks"] = data_manager.nb_tasks
    model = get_learner(args["learner_name"], args, data_manager)

    cnn_curve, nme_curve = {"top1": [], "top5": []}, {
        "top1": [],
        "top5": [],
    }
    accuracy = 0.
    for task in range(data_manager.nb_tasks):
        logging.info("All params: {}".format(count_parameters(model._network)))
        logging.info(
            "Trainable params: {}".format(count_parameters(model._network, True))
        )
        model.incremental_train(data_manager)
        cnn_accy, nme_accy = model.eval_task()
        model.after_task()

        if nme_accy is not None:
            logging.info("CNN: {}".format(cnn_accy["grouped"]))
            logging.info("NME: {}".format(nme_accy["grouped"]))

            cnn_curve["top1"].append(cnn_accy["top1"])
            cnn_curve["top5"].append(cnn_accy["top5"])

            nme_curve["top1"].append(nme_accy["top1"])
            nme_curve["top5"].append(nme_accy["top5"])

            logging.info("CNN top1 curve: {}".format(cnn_curve["top1"]))
            logging.info("CNN top5 curve: {}".format(cnn_curve["top5"]))
            logging.info("NME top1 curve: {}".format(nme_curve["top1"]))
            logging.info("NME top5 curve: {}\n".format(nme_curve["top5"]))

            print(
                "Average Accuracy (CNN):",
                sum(cnn_curve["top1"]) / len(cnn_curve["top1"]),
            )
            print(
                "Average Accuracy (NME):",
                sum(nme_curve["top1"]) / len(nme_curve["top1"]),
            )

            logging.info(
                "Average Accuracy (CNN): {}".format(
                    sum(cnn_curve["top1"]) / len(cnn_curve["top1"])
                )
            )
            logging.info(
                "Average Accuracy (NME): {}".format(
                    sum(nme_curve["top1"]) / len(nme_curve["top1"])
                )
            )
        else:
            logging.info("No NME accuracy.")
            logging.info("CNN: {}".format(cnn_accy["grouped"]))

            cnn_curve["top1"].append(cnn_accy["top1"])
            cnn_curve["top5"].append(cnn_accy["top5"])

            logging.info("CNN top1 curve: {}".format(cnn_curve["top1"]))
            logging.info("CNN top5 curve: {}\n".format(cnn_curve["top5"]))

            print(
                "Average Accuracy (CNN):",
                sum(cnn_curve["top1"]) / len(cnn_curve["top1"]),
            )
            logging.info(
                "Average Accuracy (CNN): {} \n".format(
                    sum(cnn_curve["top1"]) / len(cnn_curve["top1"])
                )
            )

        accuracy = (
            sum(cnn_curve["top1"])
            / len(cnn_curve["top1"])
            / (data_manager.nb_tasks - task)
        )
        trial.report(accuracy, task)

        if task % 3 == 2 or task == data_manager.nb_tasks - 1:
            # Handle pruning based on the intermediate value.
            if trial.should_prune():
                raise optuna.exceptions.TrialPruned()

    return accuracy


def objective(trial):
    args = setup_parser().parse_args()
    param = load_json(args.config)
    args = merge_config(args, param)
    args = NamespaceDict(**args)

    args["inc_epochs"] = trial.suggest_int("inc_epochs", 5, 20, step=5)
    args["inc_lr"] = trial.suggest_float("inc_lr", 0.0, 0.05)
    args["ca_lr"] = trial.suggest_float("ca_lr", 0.0, 0.05)

    args["distill_alpha"] = trial.suggest_float("distill_alpha", 0.0, 0.1)
    args["align_alpha"] = trial.suggest_float("align_alpha", 0.7, 0.9)

    args["margin"] = trial.suggest_float("margin", 0.0, 0.5)
    args["scale"] = trial.suggest_int("scale", 10, 30, step=5)

    obj = _optuna_train(args, trial)

    return obj


if __name__ == "__main__":
    # Add stream handler of stdout to show the messages
    optuna.logging.get_logger("optuna").addHandler(logging.StreamHandler(sys.stdout))
    study_name = "CIFAR-hp-search"  # Unique identifier of the study.
    storage_name = "sqlite:///{}.db".format(study_name)
    study = optuna.create_study(
        direction="maximize",
        pruner=optuna.pruners.HyperbandPruner(),
        sampler=optuna.samplers.TPESampler(),
        study_name=study_name,
        storage=storage_name,
        load_if_exists=True,
    )

    study.optimize(objective, n_trials=100)

    pruned_trials = study.get_trials(deepcopy=False, states=[TrialState.PRUNED])
    complete_trials = study.get_trials(deepcopy=False, states=[TrialState.COMPLETE])

    print("Study statistics: ")
    print("  Number of finished trials: ", len(study.trials))
    print("  Number of pruned trials: ", len(pruned_trials))
    print("  Number of complete trials: ", len(complete_trials))

    print("Best trial:")
    trial = study.best_trial

    print("  Value: ", trial.value)

    print("  Params: ")
    for key, value in trial.params.items():
        print("    {}: {}".format(key, value))
